import baostock as bs
import pandas as pd
import db
from datetime import datetime, timedelta
import concurrent.futures
import sys
import asyncio
import time

dateFormat = "%Y-%m-%d"
cursor_stock = ""
user_id = ""

def load_cursor():
    cursor_sql = "select * from `cursor` where table_name='bao_stock'"
    with db.UsingMysql() as um:
        um.cursor.execute(cursor_sql)
        cursor_rs = um.cursor.fetchall()
    return cursor_rs

def build_sharding(cursor_rs, sharding_rs):
    dic = {}
    for data in cursor_rs:
        sz_code = data["code"]
        sh_code = data["code"]
        dic[sz_code] = 1
        dic[sh_code] = 1
    clear_data = []
    for data in sharding_rs:
        sz_code = "sz." + data["stock_code"]
        sh_code = "sh." + data["stock_code"]
        sh_data = data.copy()
        sz_data = data.copy()
        if dic.get(sz_code) == None:
            sz_data["stock_code"] = sz_code
            clear_data.append(sz_data)
        else:
            print(sz_code, " handled")
        if dic.get(sh_code) == None:
            sh_data["stock_code"] = sh_code
            clear_data.append(sh_data)
        else:
            print(sh_code, " handled")
    return clear_data

def load_by_code(data):
    stock_code = data["stock_code"]
    index = data["sharding_table"]
    # soket = data["default_socket"]
    # user_id = data["user_id"]
    # setattr(conx, "default_socket", soket)
    # setattr(conx, "user_id", user_id)
    # start_date = datetime.strptime("2000-03-01", dateFormat)
    load_by_day(stock_code, index)


def load_by_day(stockCode, index):
    freStr = '5'
    aStr = '3'
    start_date_Str = "2000-01-01"
    rs = None
    with db.UsingMysql() as um:
        sql = "select max(`date`) as `date` from minute_line_" + index + " where code=%s"
        params = (stockCode)
        um.cursor.execute(sql, params)
        date_rs = um.cursor.fetchone()
        if date_rs["date"] != None:
            start_date = date_rs["date"]
            delta = timedelta(days=1)
            start_date = start_date - delta
            start_date_str = datetime.strftime(start_date, dateFormat)
    try:
        rs = bs.query_history_k_data(stockCode, "date,time,code,open,high,low,close,volume,amount",
        start_date=start_date_Str,frequency=freStr,adjustflag=aStr)
    except TimeoutError:
        load_by_day(stockCode, index)
    list = []
    i = 0
    tasks = []
    loop = asyncio.get_event_loop()
    with concurrent.futures.ProcessPoolExecutor(max_workers=1) as executor:
        while (rs.error_code == '0') & rs.next():
            list.append({"data":rs.get_row_data(), "index": index})
            if (i + 1) % 1000 == 0:
                tasks.append(begin_save(list, executor, None))
                list = []
            i = i+1
        if len(list) > 0:
            tasks.append(begin_save(list, executor, None))
    # asyncio.wait(tasks)
    update_cursor(stockCode)


def update_cursor(stock_code):
    sql = "insert into `cursor` (table_name, code) values(%s,%s);"
    params = ("bao_stock", stock_code)
    with db.UsingMysql() as um:
        um.cursor.execute(sql, params)

def begin_save(list, executor, call_back):
    return executor.submit(save_list, list)
    # with db.UsingMysql() as um:
    #     for data in list:
    #         if save(data, um) == False:
    #             break;

def save_list(list):
    # with concurrent.futures.ProcessPoolExecutor(50) as executor:
        # executor.map(save, list)
    with db.UsingMysql() as um:
        check_date = ""
        need_check = True
        start_time = time.time()
        for data in list:
            if check_date != data["data"][0]:
                need_check = True
            need_check = save(data, um, need_check)
            check_date = data["data"][0]
        end_time = time.time()
        print('数据量：' + str(len(list)) + '  总共的时间为:', round(end_time - start_time, 3),'secs')

def save(rs, um, need_check):
    data = rs["data"]
    index = rs["index"]
    has_repeated = False
    if need_check:
        select = "select id from minute_line_" + index + " where code=%s and time=%s"
        select_params = (data[2], data[1])
        um.cursor.execute(select, select_params);
        selectData = um.cursor.fetchone();
        has_repeated = selectData != None
    if has_repeated:
        return True
        
    sql = "insert into minute_line_" + index + " (`date`,`time`,code,open,high,low,close,volume,amount) values(%s,%s,%s,%s,%s,%s,%s,%s,%s);"
    params = (data[0],data[1],data[2],data[3],data[4],data[5],data[6],data[7],data[8])
    # try:
    um.cursor.execute(sql, params)
    um.commit()
    # except Exception:
    #     return False;
    # print(data)
    return False

def main(part, part_count):
    bs.login()
    bs.logout()

if __name__ == '__main__':
    args = ""
    if len(sys.argv) > 1:
        args = sys.argv[1]
    part = 0
    part_count = 0
    if args != "":
        argArrary = args.split('/')
        part = int(argArrary[0]) 
        part_count = int(argArrary[1]) 
    main(part, part_count)

    